from dashscope import Generation
from typing import Optional
from config import app_config
from .llm_base import BaseLLM
from .llm_registry import register_llm

from chat_memory import build_message_history, append_to_history
from ..repository.characterCard_repository import character_card_repo
from ..repository.chatSession_repository import chat_session_repo

import os

@register_llm("qwen")
class QwenLLM(BaseLLM):
    def __init__(self) -> None:
        super().__init__()  
        self._model_name = app_config.QWEN_MODEL
        # 历史记忆长度，防止超限
        self._history_limit = 5

    def generate(
        self, 
        prompt: str,
        session_id: int = -1,
        max_tokens: int = 2048,
        stream: bool = False,
        enable_memory: bool = False,  # 🔁 记忆开关
        **kwargs  # 允许扩展其他参数
    ) -> Optional[str]:
        """
        调用通义千问，支持记忆开关控制是否记住上下文

        Args:
            prompt: 用户输入文本
            session_id: 会话ID，用于区分不同用户
            model: 模型名称
            system_prompt: 系统提示词，如果会话存在角色卡，则会覆盖
            character: 角色名（如 "你是一个乐于助人的助手"）
            temperature: 温度
            top_p: 核采样
            max_tokens: 最大生成长度
            stream: 是否流式输出
            enable_memory: 是否开启记忆（记住历史）
        """
        try:
            # ===== 0. 预设参数 =====
            # 通过会话id,获取角色卡，生成系统提示词
            session_data = chat_session_repo.get_session_by_id(session_id)
            if session_data and session_data.get("character_id"):
                character_id = session_data["character_id"]
                # 根据角色卡生成系统提示词
                system_prompt = character_card_repo.build_system_prompt(character_id)
                # 根据会话读取角色卡中的指定的模型参数
                model_params = character_card_repo.get_model_params(session_id)
                if not model_params and model_params["temperature"] and model_params["top_p"]:
                    temperature = model_params["temperature"]
                    top_p = model_params["top_p"]
                else:
                    raise ValueError(f"会话 {session_id} 关联的角色卡缺少模型参数")
            else:
                raise ValueError(f"会话 {session_id} 未关联角色卡")
            # ===== 2. 构建消息历史 =====
            messages = build_message_history(
                session_id=session_id,
                system_prompt=system_prompt,
                user_input=prompt,
                enable_memory=enable_memory,
                history_limit=self._history_limit  # 根据实例属性设置历史长度
            )
            # ===== 3. 调用模型 =====
            response = Generation.call(
                api_key=os.getenv("DASHSCOPE_API_KEY"),  # type: ignore
                model=self._model_name,
                messages=messages, # type: ignore
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                result_format='message',
                stream=stream,
                incremental_output=True  # 只返回新增部分
            )
            # ===== 4. 处理响应 =====
            if not stream:
                # 非流式
                if response.status_code == 200: # type: ignore
                    reply = response.output.choices[0].message.content # type: ignore
                else:
                    raise Exception(f"qwen调用失败: {response.message}") # type: ignore
            else:
                # 流式
                reply = self._handle_streaming_response(response)
            
            # ===== 5. 保存记忆（如果启用） =====
            if enable_memory:
                append_to_history(session_id, "assistant", reply) # type: ignore

            return reply # type: ignore
        except Exception as e:
            print(f"🚨 qwen 请求异常: {str(e)}")
            return None

    def cleanup(self):
        """QwenLLM 无需显式释放资源"""
        pass

    def _handle_streaming_response(self, response) -> str:
        """处理流式响应"""
        for resp in response:
            if resp.status_code == 200:
                delta = resp.output.choices[0].message.content
                print(delta, end="", flush=True)
                reply += delta # type: ignore
            else:
                raise Exception(f"qwen流式错误: {response.message}")
        print()  # 换行
        return reply


if __name__ == "__main__":
    pass
    # 测试代码
    # llm = QwenLLM()
    # response = llm.generate("你好，介绍一下你自己吧！", session_id=1, stream=True, enable_memory=True)
    # print("最终回复：", response)